from ultralytics import YOLO

# --- 1. 模型加载与基础配置 ---
# 你可以从这里开始你的训练，选择一个预训练模型或从头开始训练。
# 预训练模型通常能加快收敛，并适用于你的数据集与 COCO 数据集有相似之处的情况。
# - 'yolov8n.pt': Nano 模型，最小最快，适合资源有限或追求速度的场景。
# - 'yolov8s.pt': Small 模型，速度和精度之间的良好平衡。
# - 'yolov8m.pt': Medium 模型，精度更高，但速度和资源消耗也更大。
# - 'yolov8l.pt': Large 模型，更高精度，更慢。
# - 'yolov8x.pt': Extra Large 模型，最高精度，最慢，资源消耗最大。
# 或者，你可以指定你之前训练好的模型路径，以便继续训练。
MODEL_PATH = 'runs/detect/YOLOv8n_iflytekUcar/base_cust_640.pt/weights/best.pt' # 建议从 yolov8n.pt 或 yolov8s.pt 开始

# 加载 YOLO 模型
# 默认是 'detect' 任务。如果你在训练分割 (segment)、姿态 (pose) 或分类 (classify) 模型，
# 请确保加载对应的模型文件（例如 'yolov8n-seg.pt' 用于分割）。
model = YOLO(MODEL_PATH)

print(f"模型 '{MODEL_PATH}' 已成功加载。")

# --- 2. 训练参数配置  ---
# 这些参数将传递给 model.train() 方法，控制训练过程。
# 你可以根据自己的数据集、硬件和目标性能进行调整。
TRAINING_PARAMS = {
    # 数据集配置文件 (data): 你的数据集 YAML 文件路径。
    # 这个文件定义了数据集的路径、类别名称和数量等信息。
    'data': 'iflytek_ucar_phone.yaml', 

    # 训练轮次 (epochs): 模型将迭代整个数据集的次数。
    # 更多的 epochs 通常会带来更好的性能，但也会增加训练时间。
    # 过多的 epochs 可能导致过拟合。
    'epochs': 150, # 建议：小数据集 50-100，大数据集 100-500，或使用 early stopping

    # 批处理大小 (batch): 每次梯度更新使用的图像数量。
    # 较大的 batch_size 可以加速训练，并可能使梯度更稳定，但会消耗更多 GPU 内存。
    # 确保你的 GPU 内存足够大，否则会报错 OOM (Out Of Memory)。
    # 如果内存不足，可以尝试减小 batch_size，或者使用 'batch=-1' 让 Ultralytics 自动选择最佳大小。
    'batch': -1, # 建议：8, 16, 32, 64 等，取决于你的 GPU 内存

    # 图像大小 (imgsz): 输入到模型的图像尺寸。
    # 可以是单个整数（如 640，表示 640x640），也可以是元组（如 (1280, 720)）。
    # - 较小的 imgsz (如 640) 训练更快，内存消耗低，但对小目标的检测能力可能下降。
    # - 较大的 imgsz (如 960, 1280) 可以保留更多细节，对小目标检测更有利，但训练更慢，内存消耗更高。
    # 尝试你 GPU 能承受的最大 imgsz。
    'imgsz': 640, # 建议：640, 960, 1280, 或 (1280, 720)

    # 设备 (device): 指定运行训练的设备。
    # 'cpu' 表示使用 CPU，0 表示使用第一个 GPU，0,1 表示使用第一个和第二个 GPU。
    # 如果你有 GPU，强烈建议设置为 GPU ID。
    'device': '0', # 'cpu' 或 '0'

    # 优化器 (optimizer): 用于更新模型权重的算法。
    # 'auto' 让 Ultralytics 自动选择 (通常是 AdamW)。
    # 'SGD' (随机梯度下降), 'Adam', 'AdamW', 'RMSProp' 等。
    'optimizer': 'auto', # 建议：'auto' 或 'AdamW'

    # 学习率 (lr0): 初始学习率。
    # 太高可能导致训练不稳定，太低可能导致收敛过慢或陷入局部最优。
    'lr0': 0.01, # 建议：0.01

    # 学习率衰减 (lrf): 最终学习率与初始学习率的比例。
    # 控制学习率在训练结束时的下降程度。
    'lrf': 0.01, # 建议：0.01 (最终学习率为初始学习率的 1%)

    # 权重衰减 (weight_decay): L2 正则化项，用于防止过拟合。
    # 较大的值会限制模型复杂度。
    'weight_decay': 0.0005, # 建议：0.0005

    # 热身 epochs (warmup_epochs): 在训练开始时，学习率会从很小的值逐渐增加到 lr0。
    # 有助于训练初期稳定性。
    'warmup_epochs': 3.0, # 建议：3.0

    # 数据加载器工作线程数 (workers): 用于加载数据的 CPU 线程数量。
    # 增加 workers 可以加速数据预处理，但会消耗更多 CPU 资源。
    # 建议设置为你的 CPU 核心数减去 1 或 2，以避免 CPU 成为瓶颈。
    'workers': 10, # 建议：根据你的 CPU 核心数调整，例如 8

    # 项目名称 (project): 训练结果保存的顶层目录。
    'project': 'runs/detect/YOLOv8n_iflytekUcar', # 建议：自定义一个有意义的名称

    # 训练名称 (name): 本次训练的子目录名称。
    # 每次运行都会在这个目录下创建一个新的文件夹，避免覆盖。
    #'name': 'custom_dataset_run_v1', # 建议：自定义一个有意义的名称，例如 'phone_detection_run_1'

    # 恢复训练 (resume): 是否从上次中断的地方恢复训练。
    # 'True' 会自动从 project/name 目录下的 last.pt 恢复。
    # 'False' 重新开始训练。
    'resume': False, # True 或 False

    # 缓存图像 (cache): 是否将图像加载到内存/显存中以加速训练。
    # 'True' (内存), 'disk' (硬盘), 'False' (不缓存)。
    # 如果数据集较小且内存充足，设置为 True 可以显著提速。
    'cache': True, # True, 'disk', 或 False

    # 预训练权重 (pretrained): 是否使用预训练的权重。
    # True 表示使用 model_path 中的权重，False 表示随机初始化权重。
    'pretrained': True, # 通常设置为 True，除非你是从头开始训练新架构

    # 验证 (val): 是否在每个 epoch 结束时运行验证。
    'val': True, # True 或 False

    # 保存最佳权重 (save_period): 每隔多少个 epoch 保存一次检查点。
    # -1 表示只保存最终和最佳模型。
    'save_period': -1, # 建议：-1 或 5 (每 5 个 epoch 保存一次)

    # 绘制结果图 (plots): 是否保存训练和验证的图表（如损失曲线、mAP 曲线等）。
    'plots': True, # True 或 False

    # 马赛克增强关闭轮次 (close_mosaic): 在训练结束前的几个 epoch 关闭 Mosaic 数据增强。
    # 有助于在训练后期稳定性能，避免过于激进的增强。
    'close_mosaic': 10, # 建议：10-20，取决于总 epochs 数量

    # 混合精度训练 (amp): 是否使用自动混合精度。
    # 可以加速训练并减少内存使用，强烈建议启用。
    'amp': True, # 建议：True
    
    # Patience (patience): 早停机制的耐心值。
    # 如果验证集上的 mAP50-95 在连续 patience 个 epochs 内没有改善，训练将停止。
    # 设置为 0 或 None 禁用早停。
    'patience': 50, # 建议：根据数据集大小和收敛速度调整，例如 50-200
    
    # 翻转增强 (fliplr): 图像水平翻转的概率。
    'fliplr': 0.5, # 建议：0.5
    
    # 裁剪增强 (crop_fraction): 图像随机裁剪的比例。
    'crop_fraction': 1.0, # 建议：1.0

    'mosaic': 1.0,       # Mosaic增强概率 (将4张图拼接，强烈建议保持1.0)
    'mixup': 0.1,        # MixUp增强概率 (图像线性混合, 建议给一个较小值如0.1)
    'copy_paste': 0.1,   # Copy-Paste增强概率 (将图中物体复制粘贴到其他图，高效模拟遮挡)
    
    # 其他高级参数 (不常用但可调)
    # 类别平衡权重 (single_cls): 如果你的数据集只有一个类别，设置为 True 可以优化。
    # 'single_cls': False,
    # 矩形训练 (rect): 是否使用矩形训练。通常在推理时使用，训练时很少用。
    # 'rect': False,
    # 余弦退火学习率调度 (cos_lr): 是否使用余弦退火学习率调度器。
    # 'cos_lr': False,
    # 模型冻结 (freeze): 冻结模型的前 N 层。可以加速训练初期，尤其是在微调时。
    # 'freeze': None, # 例如 10 (冻结前10层)
    # 多尺度训练 (multi_scale): 随机改变图像大小进行训练。可以提高泛化能力。
    'multi_scale': True,
}

# --- 3. 执行训练 ---
print("\n--- 正在开始模型训练 ---")
print(f"训练参数概览: {TRAINING_PARAMS}")

try:
    # 启动训练
    results = model.train(**TRAINING_PARAMS)
    print("\n模型训练完成！")
    print(f"训练结果保存在: {model.trainer.save_dir}")

except Exception as e:
    print(f"\n训练过程中发生错误: {e}")
    print("请检查你的参数设置、数据集路径和 GPU 内存。")

print("\n--- 脚本执行完毕 ---")